package com.campus.counseling.service.impl;

import com.campus.counseling.service.LSTMModelTrainer;
import com.campus.counseling.service.ModelEvaluationService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nd4j.linalg.api.buffer.DataType;
import org.springframework.stereotype.Service;
import com.campus.counseling.model.ModelMetrics;
import org.nd4j.linalg.ops.transforms.Transforms;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
@RequiredArgsConstructor
public class ModelEvaluationServiceImpl implements ModelEvaluationService {

    private final LSTMModelTrainer modelTrainer;

    @Override
    public ModelMetrics evaluateModel(DataSet testData) {
        try {
            INDArray predictions = modelTrainer.predict(testData.getFeatures());
            INDArray labels = testData.getLabels();
            
            // 记录一些预测样例
            log.info("预测样例对比:");
            for (int i = 0; i < Math.min(5, predictions.size(0)); i++) {
                double pred = predictions.getDouble(i);
                double actual = labels.getDouble(i);
                log.info("样例 {}: 预测值={}, 实际值={}, 差异={}", 
                    i, String.format("%.4f", pred), 
                    String.format("%.4f", actual), 
                    String.format("%.4f", Math.abs(pred - actual)));
            }
            
            // 计算各项指标
            double mse = calculateMSE(predictions, labels);
            double mae = calculateMAE(predictions, labels);
            double accuracy = calculateAccuracy(predictions, labels);
            double precision = calculatePrecision(predictions, labels);
            double recall = calculateRecall(predictions, labels);
            double f1Score = calculateF1Score(precision, recall);
            
            // 记录详细的评估结果
            log.info("详细评估指标:");
            log.info("MSE: {}", String.format("%.4f", mse));
            log.info("MAE: {}", String.format("%.4f", mae));
            log.info("准确率: {}%", String.format("%.2f", accuracy * 100));
            log.info("精确率: {}%", String.format("%.2f", precision * 100));
            log.info("召回率: {}%", String.format("%.2f", recall * 100));
            log.info("F1分数: {}%", String.format("%.2f", f1Score * 100));
            
            return ModelMetrics.builder()
                .mse(mse)
                .mae(mae)
                .accuracy(accuracy)
                .precision(precision)
                .recall(recall)
                .f1Score(f1Score)
                .build();
            
        } catch (Exception e) {
            log.error("模型评估失败: ", e);
            throw new RuntimeException("模型评估失败", e);
        }
    }

    @Override
    public double crossValidate(List<DataSet> folds) {
        double totalScore = 0;
        int k = folds.size();
        
        for (int i = 0; i < k; i++) {
            // 准备训练集和验证集
            DataSet validationSet = folds.get(i);
            DataSet trainingSet = mergeDataSets(folds, i);
            
            // 训练模型
            modelTrainer.initModel();
            modelTrainer.train(trainingSet);
            
            // 评估当前折
            ModelMetrics metrics = evaluateModel(validationSet);
            totalScore += metrics.getF1Score();
            
            log.info("第{}折交叉验证完成，F1得分: {}", i + 1, metrics.getF1Score());
        }
        
        double averageScore = totalScore / k;
        log.info("{}折交叉验证平均F1得分: {}", k, averageScore);
        
        return averageScore;
    }

    private DataSet mergeDataSets(List<DataSet> folds, int excludeIndex) {
        DataSet merged = null;
        
        for (int i = 0; i < folds.size(); i++) {
            if (i != excludeIndex) {
                if (merged == null) {
                    merged = folds.get(i).copy();
                } else {
                    merged.merge((List<? extends org.nd4j.linalg.dataset.api.DataSet>) folds.get(i));
                }
            }
        }
        
        return merged;
    }

    private double calculateAccuracy(INDArray predictions, INDArray labels) {
        try {
            // 先转换为浮点类型再进行比较
            INDArray predFloat = predictions.castTo(DataType.FLOAT);
            INDArray labelsFloat = labels.castTo(DataType.FLOAT);
            
            INDArray binaryPredictions = predFloat.gt(0.5).castTo(DataType.FLOAT);
            INDArray binaryLabels = labelsFloat.gt(0.5).castTo(DataType.FLOAT);
            
            return binaryPredictions.eq(binaryLabels).castTo(DataType.FLOAT).meanNumber().doubleValue();
        } catch (Exception e) {
            log.error("计算准确率失败: {}", e.getMessage());
            return 0.0;
        }
    }

    private double calculatePrecision(INDArray predictions, INDArray labels) {
        try {
            // 确保数据是2维的并且维度匹配
            if (predictions.rank() == 3) {
                predictions = predictions.get(NDArrayIndex.all(), NDArrayIndex.point(0), NDArrayIndex.point(0));
            }
            if (labels.rank() == 3) {
                labels = labels.get(NDArrayIndex.all(), NDArrayIndex.point(0), NDArrayIndex.point(0));
            }
            
            // 确保维度一致
            if (labels.rank() == 1) {
                labels = labels.reshape(labels.length(), 1);
            }
            if (predictions.rank() == 1) {
                predictions = predictions.reshape(predictions.length(), 1);
            }

            // 转换为二分类结果
            // 对于情感分数，我们可以设定阈值，比如0.6表示积极情感
            double threshold = 0.6;  // 可以根据实际需求调整
            INDArray binaryPredictions = predictions.gt(threshold).castTo(DataType.FLOAT);
            INDArray binaryLabels = labels.gt(threshold).castTo(DataType.FLOAT);
            
            // 计算真阳性和预测阳性总数
            INDArray truePositives = binaryPredictions.mul(binaryLabels);
            double tp = truePositives.sumNumber().doubleValue();
            double totalPredictedPositive = binaryPredictions.sumNumber().doubleValue();
            
            // 避免除以0
            return totalPredictedPositive > 0 ? tp / totalPredictedPositive : 0.0;
        } catch (Exception e) {
            log.error("计算精确率失败: {}", e.getMessage());
            return 0.0;
        }
    }

    private double calculateRecall(INDArray predictions, INDArray labels) {
        try {
            // 确保数据是2维的并且维度匹配
            if (predictions.rank() == 3) {
                predictions = predictions.get(NDArrayIndex.all(), NDArrayIndex.point(0), NDArrayIndex.point(0));
            }
            if (labels.rank() == 3) {
                labels = labels.get(NDArrayIndex.all(), NDArrayIndex.point(0), NDArrayIndex.point(0));
            }
            
            // 确保维度一致
            if (labels.rank() == 1) {
                labels = labels.reshape(labels.length(), 1);
            }
            if (predictions.rank() == 1) {
                predictions = predictions.reshape(predictions.length(), 1);
            }

            // 使用相同的阈值
            double threshold = 0.6;
            INDArray binaryPredictions = predictions.gt(threshold).castTo(DataType.FLOAT);
            INDArray binaryLabels = labels.gt(threshold).castTo(DataType.FLOAT);
            
            // 计算真阳性和实际阳性总数
            INDArray truePositives = binaryPredictions.mul(binaryLabels);
            double tp = truePositives.sumNumber().doubleValue();
            double totalActualPositive = binaryLabels.sumNumber().doubleValue();
            
            return totalActualPositive > 0 ? tp / totalActualPositive : 0.0;
        } catch (Exception e) {
            log.error("计算召回率失败: {}", e.getMessage());
            return 0.0;
        }
    }

    private double calculateF1Score(double precision, double recall) {
        if (precision + recall == 0) {
            return 0.0;
        }
        return 2 * (precision * recall) / (precision + recall);
    }

    private double calculateMSE(INDArray predictions, INDArray labels) {
        INDArray diff = predictions.sub(labels);
        return diff.mul(diff).meanNumber().doubleValue();
    }

    private double calculateMAE(INDArray predictions, INDArray labels) {
        return Transforms.abs(predictions.sub(labels)).meanNumber().doubleValue();
    }

    private void logEvaluationExplanation(ModelMetrics metrics) {
        log.info("评估结果解释:");
        log.info("1. MSE(均方误差) = {}: 预测值与实际值的平均平方差，越小越好", 
            String.format("%.4f", metrics.getMse()));
        log.info("2. 准确率 = {}%: 所有预测中正确的比例", 
            String.format("%.2f", metrics.getAccuracy() * 100));
        log.info("3. 精确率 = {}%: 预测为积极情感(>0.6)中实际正确的比例", 
            String.format("%.2f", metrics.getPrecision() * 100));
        log.info("4. 召回率 = {}%: 实际积极情感中被正确预测的比例", 
            String.format("%.2f", metrics.getRecall() * 100));
        log.info("5. F1分数 = {}%: 精确率和召回率的调和平均，综合评价指标", 
            String.format("%.2f", metrics.getF1Score() * 100));
    }
} 